import numpy as np
import matplotlib.pyplot as plt
import scipy

def dct2(x):
    return scipy.fftpack.dct(scipy.fftpack.dct(x.T, norm='ortho').T, norm='ortho')

# 二维离散余弦反变换
def idct2(x):
    return scipy.fftpack.idct(scipy.fftpack.idct(x.T, norm='ortho').T, norm='ortho')

# 参数设置
Nx = 32
Ny = 32

# 网格间距
dx = 2 / Nx
dy = 2 / Ny

# 创建网格
x = np.arange(-1 + dx / 2, 1 - dx / 2 + dx/10, dx)
y = x = np.arange(-1 + dy / 2, 1 - dy / 2 + dy/10, dy)

X, Y = np.meshgrid(x, y, indexing='ij')  # 使用 'ij' 索引以匹配 MATLAB 的行为

# 修改波数
kx = np.arange(Nx)
ky = np.arange(Ny)
mwx = 2 * (np.cos(np.pi * kx / Nx) - 1) / dx**2
mwy = 2 * (np.cos(np.pi * ky / Ny) - 1) / dy**2

# 2D DCT of f (Right hand side)
f = -2 * np.pi**2 * np.cos(np.pi * X) * np.cos(np.pi * Y)
analy= np.cos(np.pi * X) * np.cos(np.pi * Y)

plt.figure(figsize=(12, 6))

# plt.subplot(1, 2, 1)
# plt.title("Numerical Solution")
# plt.imshow(np.real(f), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
# plt.colorbar(label='Pressure')

fhat = dct2(f)  # 使用 DCT-II 类型
# f_f=idct2(fhat)
# print(np.max(fhat))

# plt.subplot(1, 2, 2)
# plt.title("Numerical Solution")
# plt.imshow(np.real(f_f), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
# plt.colorbar(label='Pressure')

# 构造 MWX 和 MWY 的网格
MWX, MWY = np.meshgrid(mwx, mwy, indexing='ij')

# 解方程
uhat = fhat / (MWX + MWY)

# 由于 uhat(0,0) 可能未定义，我们将其设置为 0
uhat[0, 0] = 0

# 逆 2D DCT
u = idct2(uhat)

# 绘制结果




plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.title("Numerical Solution")
plt.imshow(np.real(u), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

plt.subplot(1, 2, 2)
plt.title("Analytical Solution")
plt.imshow(np.real(analy), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

plt.tight_layout()
plt.show()